"""
TrademarkScraper — 使用 Playwright 异步 API 的商标公告爬虫类
- 初始化时传入账号、密码，公告期号（可选）
- 自动为每个账号创建/使用 user_data_dir（属性：user_data_dir）并把账号保存到实例属性
- 提供登录检测方法 is_logged_in()
- 提供 login() 登录方法
- 提供 fetch_list() 爬取列表方法（示例返回字典列表）
- 能检测 301/302 重定向（在 response 事件中检查状态码）

注意：此代码为模板，需要根据目标网站的具体页面结构（选择器、登录流程）做少量修改。
"""

from typing import Optional, List, Dict
from playwright.async_api import async_playwright, Browser, BrowserContext, Page, Response
from mytools import *
from sql import MySQLDB
import os,asyncio,json,traceback,math


class TrademarkScraper:
    def __init__(self, username: str, password: str, gonggaoqi: Optional[int] = None, bfs=1, user_data_dir: Optional[str] = None,db_list:List[MySQLDB]=None,table_name='zgsb',table_name2='zgsbqs'):
        """
        username, password: 账号密码
        gonggaoqi: 公告期号，可选
        user_data_dir: 可指定持久化目录；如果不指定，会自动生成 user_data_{username}
        """
        self.username = username
        self.password = password
        self.gonggaoqi = gonggaoqi
        self.bfs=bfs
        self.db_list=db_list
        self.db=db_list[0]
        self.table_name=table_name
        self.table_name2=table_name2
        self.save_qq_status=False
        self.sep_qq=asyncio.Semaphore(1)
        self.stop_cj_plist=False
        
        # 自动分配用户目录（以账号安全命名）
        if user_data_dir:
            self.user_data_dir = user_data_dir
        else:
            safe_name = ''.join(c for c in username if c.isalnum() or c in ('-', '_'))
            self.user_data_dir = os.path.abspath(f"user_data_{safe_name}")
            os.makedirs(self.user_data_dir, exist_ok=True)

        # Playwright 资源占位
        self._playwright = None
        self.browser: Optional[Browser] = None
        self.context: Optional[BrowserContext] = None
        self.page: Optional[Page] = None

        # 可以在实例上保存最后一次抓取的 cookies/session（示例）
        self.last_session_file = os.path.join(self.user_data_dir, 'session.json')

        # 存放捕获到的重定向记录
        self.redirects: List[Dict] = []
        self.yzm_code=None


    async def start(self, headless: bool = False, extra_args: Optional[List[str]] = None):
        """启动 Playwright 和持久化上下文（用于维持登录状态）"""
        self._playwright = await async_playwright().start()
        args = extra_args or [
            "--start-maximized",
            "--disable-infobars",
            "--disable-blink-features=AutomationControlled",
            "--window-size=1920,1080"
        ]
        self.browser = await self._playwright.chromium.launch_persistent_context(
            user_data_dir=self.user_data_dir,
            headless=headless,
            viewport=None,
            args=args,
            accept_downloads=True
        )
        # launch_persistent_context 返回的是 BrowserContext
        self.context = self.browser
        # 取第一个页面或 new_page
        self.page = self.context.pages[0] if self.context.pages else await self.context.new_page()

        # 注册 response 监听，检测 3xx 重定向
        #self.page.on('response', self._on_response)

    async def _on_response(self, response: Response):
        try:
            status = response.status
            url = response.url
            
            if status in (301, 302, 303, 307, 308):
                self.redirects.append({'status': status, 'url': url, 'headers': dict(response.headers)})
            
            if '/cas/captcha.jpg' in url:
                content: bytes = await response.body()
                cur_yzm_code=get_yzm(content)
                self.yzm_code=cur_yzm_code
                print(f'{self.username},捕获到登录验证码:{cur_yzm_code}')

        except Exception:
            # 忽略监听中的任意异常
            return

    async def stop(self, close_context: bool = False):
        """
        结束 Playwright。对于持久化上下文：如果你要退出并关闭浏览器进程，可以把 close_context=True
        - close_context=False 时只会停止 playwright 运行时（context 保持），通常不推荐这样做。
        - 对于 launch_persistent_context，close() 会关闭浏览器窗口并保存用户数据
        """
        if self.context and close_context:
            await self.context.close()
            self.context = None
            self.page = None
        if self._playwright:
            await self._playwright.stop()
            self._playwright = None

    async def is_logged_in(self,page:Page) -> bool:
        """
        检测是否登录：示例实现为访问个人中心或判断页面 DOM 中是否存在登录后才有的元素
        需要根据目标站点自定义判断逻辑
        """
        if not page:
            raise RuntimeError('请先调用 start() 启动浏览器')

        try:
            el = await page.query_selector('xpath=//div[text()="商标公告列表"]')
            if el:
                #print(f'{self.username},已登录')
                return True
        except Exception:
            pass

        # 最保险：检查是否出现登录表单
        login_form = await page.query_selector('form#loginForm')
        return login_form is None

    async def login(self,page:Page,jtq=None) -> bool:
        """
        执行登录操作。此处为模板示例：打开登录页、填用户名密码、提交表单
        返回 True 表示登录成功
        """
        
        print(f'账号:《{self.username}》,未登录进行登录')
        await asyncio.sleep(3)
        await page.fill('input[name="username"]', self.username)
        await asyncio.sleep(0.5)
        await page.fill('input[name="password"]', self.password)

        #输入验证码
        yzm_code=None
        if jtq:
            try:
                async with jtq as resp_info:
                    resp=await resp_info.value
                    content=await resp.body()
                    yzm_code=get_yzm(content)
                    print(f'{self.username},捕获到登录验证码:{yzm_code}')
            except Exception as e:
                print(f'登录页面监听验证码错误 => {e}')
                
        while True:
            try:
                el_yz_ipt= page.locator('input#captcha')
                if (await el_yz_ipt.count())==0:
                    break
                yz_value= await el_yz_ipt.input_value()
                if yz_value and len(yz_value)>=4:
                    print(f'验证码输入完毕,当前输入验证码:{yz_value}')
                    break
                el_yz_img=page.locator('img#captcha_img')
                if (await el_yz_img.count())>0 and yzm_code:
                    await el_yz_ipt.fill(yzm_code)

            except Exception as e:
                traceback.print_exc()
                print(e)
            await asyncio.sleep(3)
        
        el_check=page.locator('input[name="rememberMe"]')
        for iii in range (await el_check.count()):
            try:
                checked = await el_check.nth(iii).is_checked()
                if (await el_check.nth(iii).is_visible()) and not checked:
                    await el_check.nth(iii).check()
                    await asyncio.sleep(1)
            except Exception as e:
                traceback.print_exc()
                print(e)

        for _ in range(3):
            try:
                await page.click('input[id="btn"]')
                break
            except Exception as e:
                traceback.print_exc()
                print(e)

        # 等待跳转完成或某个登录后可见元素出现
        
        for _ in range(3):
            try:
                await page.wait_for_load_state('networkidle',timeout=90000)
                login_form = await page.query_selector('form#login')
                break
            except Exception as e:
                traceback.print_exc()
                print(e)
            await asyncio.sleep(3)
        success = login_form is None

        if success:
            try:
                cookies = await self.context.cookies()
                with open(self.last_session_file, 'w', encoding='utf-8') as f:
                    json.dump(cookies, f, ensure_ascii=False, indent=2)
            except Exception:
                pass

        return success
    
    async def plw_bring_to_front_if_needed(self,page: Page):
        """只在页面不在前台时 bringToFront，后台循环检测"""
        while True:
            try:
                is_focused = await page.evaluate("document.hasFocus()")
                if not is_focused:
                    await page.bring_to_front()
            except Exception:
                pass  # 可以加日志记录错误
            await asyncio.sleep(random.randint(30,40))

    async def check_ggqs(self,page:Page,pj):

        el_gonggaoqi=page.locator('input#basic_anncIssue')
        el_cht_btn=page.locator('xpath=//span[text()="查 询"]')
        for _ in range(5):
            try:
                if (await el_gonggaoqi.count())>0:
                    cur_ggq=await el_gonggaoqi.input_value()
                    cur_ggq=int(cur_ggq.strip())
                    if not self.gonggaoqi:
                        self.gonggaoqi=cur_ggq
                        print(f'账号:《{self.username}》,第 {pj} 个标签,获取公告期数:{self.gonggaoqi}')
                        break
                    else:
                        if self.gonggaoqi==cur_ggq:
                            print(f'账号:《{self.username}》,第 {pj} 个标签,无需更改公告期数')
                        else:
                            await el_gonggaoqi.fill(str(self.gonggaoqi))
                            await asyncio.sleep(0.5)
                            await el_cht_btn.click(delay=200)
                            await asyncio.sleep(5)
                            cur_ggq=await el_gonggaoqi.input_value()
                            cur_ggq=int(cur_ggq.strip())
                            if cur_ggq ==self.gonggaoqi:
                                print(f'账号:《{self.username}》,第 {pj} 个标签,成功切换公告期数置:{cur_ggq}')
                                break
                
            except Exception as e:
                traceback.print_exc()
                print(e)
            await asyncio.sleep(1)

    async def jump_page(self,pj,page:Page,target_p):

        print(f'账号:《{self.username}》,第 {pj} 个标签,跳转到第 {target_p} 页')
        try:
            el_jump_ipt=page.locator('xpath=//div[@class="ant-pagination-options-quick-jumper"]/input[1]')
            await el_jump_ipt.fill(str(target_p))
            await asyncio.sleep(1)
            await el_jump_ipt.press('Enter')
            await asyncio.sleep(1)
            return True
        except Exception as e:
            print(f'账号:《{self.username}》,第 {pj} 个标签,跳转到 {target_p} 页失败,正在重试...')
            return False


    async def fetch_list(self,pj,page:Page,tlist) -> Dict:
        """
        爬取商标公告列表的示例方法
        - page_no: 页码
        返回：简单的字典列表
        注意：需要根据目标站点的实际 HTML 结构实现解析
        """
        min_p,max_p=tlist[0],tlist[-1]
        print(f'账号:《{self.username}》,第 {pj} 个标签,需采集页数范围:{min_p}~{max_p},共 {len(tlist)} 页')
        await asyncio.sleep((pj-1)*60)
        t0=asyncio.create_task(self.plw_bring_to_front_if_needed(page))
        api_pd='/toas-pub-prod/pub-prod-api/public/web/anncInfo/searchEsTmgg'
        
        return_res={'code':0,'data':None,'pj':pj,'err_msg':None}
        
        need_do_ind=1
        page_no=None
        for jjj,p in enumerate(tlist,start=1):
            try:
                if self.stop_cj_plist:
                    print(f'账号:《{self.username}》,第 {pj} 个标签,停止采集')
                    break
                jtq_yzimg=page.expect_response(lambda r: '/cas/captcha.jpg' in r.url and r.status == 200,timeout=90000)
                if jjj==need_do_ind:
                    list_url = f'https://cas.sbj.cnipa.gov.cn/cas/login?service=https://pub.sbj.cnipa.gov.cn/toas-pub-prod/portalui-pub-prod/brandNotice'
                    await page.goto(list_url, wait_until='networkidle',timeout=90000)
                    await asyncio.sleep(5)

                    logged = await self.is_logged_in(page)
                    if not logged:
                        ok = await self.login(page,jtq_yzimg)
                        print(f'账号:《{self.username}》,第 {pj} 个标签,登录结果:{ok}')
                        if not ok:
                            continue
                        await asyncio.sleep(5)

                    await self.check_ggqs(page,pj) #检查公告期数
                    if p==1:
                        jtq=page.expect_response(lambda r: api_pd in r.url and r.status == 200,timeout=90000)
                    
                    wait_psize=0
                    el_pszie= page.locator('xpath=//li[@class="ant-pagination-options"]/div[1]')
                    is_tc=False
                    while True:
                        try:
                            await el_pszie.click(delay=200)
                            await asyncio.sleep(3)
                            wait_psize+=3
                            el_max_pasize=page.locator('xpath=//div[@class="rc-virtual-list-holder-inner"]/div[4]')
                            if (await el_max_pasize.count()>0):
                                break
                        except Exception as e:
                            print(e)
                        
                        await asyncio.sleep(1)
                        wait_psize+=1
                        if wait_psize>30:
                            is_tc=True
                            break
                    if is_tc:
                        print(f'账号:《{self.username}》,第 {pj} 个标签,第 {p} 页,切换page_size失败,正在重试...')
                        need_do_ind+=1
                        continue

                    print(f'账号:《{self.username}》,第 {pj} 个标签,成功更改每页数量为100个')
                    await el_max_pasize.click(delay=200)
                    await asyncio.sleep(1)


                    if jjj==need_do_ind and p!=1:
                        jtq=page.expect_response(lambda r: api_pd in r.url and r.status in [200,403],timeout=90000)
                        is_suc_jump=await self.jump_page(pj,page,p)
                        if not is_suc_jump:
                            continue

                else:
                    logged = await self.is_logged_in(page)
                    if not logged:
                        ok = await self.login(page,jtq_yzimg)
                        print(f'账号:《{self.username}》,第 {pj} 个标签,登录结果:{ok}')
                        if not ok:
                            continue
                        await asyncio.sleep(5)
                    jtq=page.expect_response(lambda r: api_pd in r.url and r.status == 200,timeout=90000)
                    if page_no and (page_no+1)==p:
                        el_next_btn=page.locator('xpath=//li[@class="ant-pagination-next"]')
                        for _ in range(5):
                            try:
                                if (await el_next_btn.is_visible()):
                                    await el_next_btn.click(delay=200)
                                    break
                            except Exception as e:
                                print(e)
                            await asyncio.sleep(3)
                    else:
                        await self.jump_page(pj,page,p)


                is_403=False
                try:
                    async with jtq as resp_info:
                        
                        resp=await resp_info.value
                        if resp.status==200:
                            json_data=await resp.json()
                            
                            plist=json_data['data']['list']
                            page_no=json_data['data']['pageIndex']
                            all_ps=json_data['data']['pages']
                            pzz=json_data['data']['pageSize']
                            
                            # with open(f'测试数据/第{page_no}页.json','w',encoding='utf-8') as f:
                            #     json.dump(json_data,f,ensure_ascii=False)
                            if page_no==1:
                                print(f'期数:{self.gonggaoqi},还未采集过,先存入所有页数...')
                                aff_qq,aff_qq_has=await self.yb_save_ggqs(all_ps,pzz)
                                print(f'期数:{self.gonggaoqi},还未采集过,成功存入 {aff_qq} 个页数,已存在忽略:{aff_qq_has}')

                            aff,aff_has,aff_cg=await self.yb_save_plist_to_mysql(pj,plist,page_no)
                            print(f'账号:《{self.username}》,第 {pj} 个标签,成功截取第 {page_no} 页内容,{len(plist)} 个,保存:{aff},已存在忽略:{aff_has},更改采集状态:{aff_cg}')

                            print(f'账号:《{self.username}》,第 {pj} 个标签,剩余采集页数:{max_p-page_no}')
              

                            if page_no==max_p:
                                print(f'账号:《{self.username}》,第 {pj} 个标签,采集至页数 {page_no},达到最大页数:{max_p},退出采集')
                                break
                        elif resp.status==403:
                            is_403=True
                            print(f'账号:《{self.username}》,账号403,请切换IP,或者等待5分钟...')

                    if is_403:
                        return_res['code']=403
                        return_res['err_msg']='账号403,请切换IP,或者等待5分钟'
                        await asyncio.sleep(300)
                        
                except Exception as e:
                    traceback.print_exc()
                    print(f'账号:《{self.username}》,第 {pj} 个标签,第 {p} 页等待截取接口出错 => {e}')
            except Exception as e:
                traceback.print_exc()
                print(f'账号:《{self.username}》,第 {pj} 个标签,第 {p} 页出错 => {e}')
            await yb_human_wait()
        await page.close()
        return return_res

    async def cj_from_mutul_page(self):
        tb2=self.table_name2
        ggqs=self.gonggaoqi
        rs=self.db.query(f'select page,state from {tb2} where page_size=100 and anncIssue=%s order by page',(str(ggqs),))
        if len(rs)==0:
            task_list=list(range(1,1001))
            print(f'公告期数:{ggqs},首次采集')
        else:
            task_list=[p for p,s in rs if s==0]

        
        if len(task_list)==0:
            print(f'公告期数:{ggqs},已全部采集完毕')
            return {'code':0,'msg':'已全部采集完毕'}

        if not self.context:
            await self.start()

        yb_tasks=[]
        task_len=math.ceil(len(task_list)/self.bfs)
        for pj in range(self.bfs):
            cur_task_list=task_list[pj*task_len:(pj+1)*task_len]
            page=await self.context.new_page()
            yb_tasks.append(asyncio.create_task(self.fetch_list(pj+1,page,cur_task_list)))
        
        await asyncio.gather(*yb_tasks)

        await self.stop()

    def save_plist_to_mysql(self,pj,plist,ppp):
        try:
            db=self.db_list[pj-1]
            tb_name=self.table_name
            tb2=self.table_name2
            aff=0
            aff_has=0
            cur_qs=self.gonggaoqi
            for pdata in plist:
                pdata['_id']=pdata.pop('id')
                pdata['imgDir']=f"https://sbggwj.sbj.cnipa.gov.cn:9443/tmann{pdata['imgDir']}"
                rs= db.query(f'select count(*) from {tb_name} where regNo=%s and anncIssue=%s',(pdata['regNo'],pdata['anncIssue']))
                if rs[0][0]==0:
                    zds=[]
                    paras=[]
                    for k,v in pdata.items():
                        zds.append(k)
                        if isinstance(v,list):
                            paras.append('|'.join(v))
                        else:
                            paras.append(v)
                    zdstr=','.join(zds)
                    wstr=','.join([r'%s' for _ in paras])
                    sqlstr=f'insert into {tb_name}({zdstr}) values({wstr})'
                    aff+=db.execute_command(sqlstr,paras)
                else:
                    aff_has+=1
            if aff or aff_has:
                add_cg=db.execute_command(f'update {tb2} set state=1 where anncIssue=%s and page=%s and page_size=100',(str(cur_qs),ppp))
            return aff,aff_has,add_cg
        except Exception as e:
            traceback.print_exc()
            print(e)
            return 0, 0,0


    async def yb_save_plist_to_mysql(self,pj,plist,cpage):
        loop = asyncio.get_running_loop()
        aff,aff_has,aff_cg = await loop.run_in_executor(None, self.save_plist_to_mysql,pj,plist,cpage)
        return aff,aff_has,aff_cg
    
    async def yb_save_ggqs(self,all_pages,psize):
        async with self.sep_qq:
            if self.save_qq_status:
                return
            self.save_qq_status=True
            ggqs=self.gonggaoqi
            tb2=self.table_name2
            def save_ggqs():
                aff=0
                aff_has=0
                print(f'期数:{ggqs},共:{all_pages} 页')
                for ppp in range(1,all_pages+1):
                    rs=self.db.query(f'select count(*) from {tb2} where anncIssue=%s and page=%s and page_size=%s',(str(ggqs),ppp,psize))
                    if rs[0][0]==0:
                        pdata={
                            'anncIssue':str(ggqs),
                            'page':ppp,
                            'page_size':psize,
                            'max_page':all_pages
                        }
                        zdstr=','.join(pdata.keys())
                        paras=list(pdata.values())
                        wstr=','.join([r'%s' for _ in pdata.values()])
                        sqlstr=f'insert into {tb2}({zdstr}) values({wstr})'
                        aff+=self.db.execute_command(sqlstr,paras)
                    else:
                        aff_has+=1

                return aff,aff_has
            
            loop = asyncio.get_running_loop()
            aff,aff_has = await loop.run_in_executor(None, save_ggqs)
            self.save_qq_status=False
            return aff,aff_has

    async def yb_do_sql(self,sql,paras):
        
        def do_sql(csql,cparas):
            aff=self.db.execute_command(csql,cparas)
            return aff
        loop = asyncio.get_running_loop()
        res_aff= await loop.run_in_executor(None, do_sql,sql,paras)
        return res_aff

    def start_yb_task(self):
        asyncio.run(self.cj_from_mutul_page())

# if __name__ == '__main__':

#     db_list=[MySQLDB('tiaoma') for j in range(5)]
#     s = TrademarkScraper('maoge55', 'Bianma001!',1960,2,db_list=db_list)
#     s.start_yb_task()
